// Created: Dec 14, 2014
//
// Author: wjinwen.1988@gmail.com

#include "chi/net/tls_connection.h"

#include <openssl/err.h>
#include <glog/logging.h>

#include "chi/net/channel.h"
#include "chi/net/event_loop.h"
#include "chi/net/socket.h"

namespace chi {
namespace net {

// This converts 'illegal' shutdowns int ZERO_RETURN.
inline bool zero_return(int error, int rc) {
  return (error == SSL_ERROR_ZERO_RETURN || (rc == 0 && errno == 0));
}

TLSConnection::TLSConnection(SSL* ssl,
                             EventLoop* loop,
                             const std::string& name,
                             int sockfd, 
                             const InetAddress& local_addr,
                             const InetAddress& peer_addr) 
  : TcpConnection(loop, name, sockfd, local_addr, peer_addr),
    ssl_state_(SSLStateEnum::STATE_UNINIT),
    ssl_(ssl) {
  VLOG(11) << "TLSConnection::ctor[" << name_ << "] at fd=" << sockfd;
}

TLSConnection::~TLSConnection() {
  VLOG(11) << "TLSConnection::dtor[" << name_ << "] at fd=" << channel_->fd();
  
  if (ssl_) {
    SSL_shutdown(ssl_);
    SSL_free(ssl_);
  } 
}

void TLSConnection::Send(const std::string& message) {
  Send(message.data(), message.size());
}

void TLSConnection::Send(const void* data, size_t len) {
  if (state_ == StateE::kConnected) {
    std::string message(static_cast<const char*>(data), len);
    loop_->RunInLoop(std::bind(&TLSConnection::SendInLoop, this, message));
  }
}

void TLSConnection::Send(Buffer& message) {
  Send(message.Peek(), message.ReadableBytes());
}

void TLSConnection::SSLAccept() {
  if (ssl_state_ != SSLStateEnum::STATE_UNINIT) {
    LOG(ERROR) << "SSLAccept() called in invalid state";
    return ;
  }

  state_ = StateE::kConnected;
  ssl_state_ = SSLStateEnum::STATE_ACCEPTING;

  // enable reading (waiting for CLIENT HELLO)
  channel_->EnableReading();
}

void TLSConnection::HandleRead() {
  if (state_ < StateE::kConnected) {
    return TLSConnection::PerformRead();
  }

  if (ssl_state_ == SSLStateEnum::STATE_ACCEPTING) {
    HandleAccept();
    return;
  }

  TLSConnection::PerformRead();
}

void TLSConnection::HandleWrite() {
  VLOG(11) << "TLSConnection::HandleWrite() this=" << this << " fd=" << socket_->fd() 
           << ", state=" << int(state_) << ", ssl_state=" << int(ssl_state_) 
           << ",events=" << channel_->events();
  if(state_ < StateE::kConnected) {
    return TLSConnection::PerformWrite();
  }

  if (ssl_state_ == SSLStateEnum::STATE_ACCEPTING) {
    HandleAccept();
    return;
  } 

  TLSConnection::PerformWrite();
}

void TLSConnection::HandleClose() {
  VLOG(11) << "fd=" << channel_->fd() << " state=" << state_;
  assert(state_ == StateE::kConnected || state_ == StateE::kDisconnecting);
  set_state(StateE::kDisconnected);
  channel_->DisableAll();
  SSL_shutdown(ssl_);

  connection_callback_(shared_from_this());
  close_callback_(shared_from_this());
}

void TLSConnection::HandleError() {
  LOG(ERROR) << "socket error";
}

void TLSConnection::SendInLoop(const std::string& message) {
  RealSend(message.data(), message.size());
}

void TLSConnection::RealSend(const void* data, size_t len) {
  ssize_t nwrote = 0;
  ssize_t remaining = len;
  bool error = false;
  if (state_ == StateE::kDisconnected || 
      ssl_state_ != SSLStateEnum::STATE_ESTABLISHED) {
    LOG(WARNING) << "disconnected, give up writing";
    return;
  }
  // if nothing in output queue, try writing directly
  if (!channel_->IsWriting() && output_buffer_.ReadableBytes() == 0) {
    nwrote = SSL_write(ssl_, data, len);
    if (nwrote <= 0) {
      long ssl_error = SSL_get_error(ssl_, nwrote);
      if (ssl_error == SSL_ERROR_WANT_WRITE) {
        // no data writtern.
        ERR_clear_error();
      } else if (ssl_error == SSL_ERROR_WANT_READ) {
        LOG(ERROR) << "TLSConnection(fd=" << socket_->fd() << ", state=" 
                   << int(state_) << ", ssl_state=" << int(ssl_state_)
                   << "events=" << channel_->events() << "):"
                   << "unspported SSL renegotiation during write";
        ERR_clear_error();
        error = true;
      } else {
        long last_error = ERR_get_error();
        LOG(ERROR) << "TLSConnection(fd=" << socket_->fd() << ", state="
                   << int(state_) << ", ssl_state_=" << int(ssl_state_) 
                   << ", events=" << channel_->events() << "): SSL error: "
                   << error << ", errno: " << errno << ", func: " 
                   << ERR_func_error_string(last_error) << ", reason: "
                   << ERR_reason_error_string(last_error);
        ERR_clear_error();
      }
    } else {
      remaining = len - nwrote;
    }
  }

  assert((size_t)remaining <= len);
  if (!error && remaining > 0) {
    VLOG(11) << "I am going to write more data";
    output_buffer_.Append(static_cast<const char*>(data)+nwrote, remaining);
    if (!channel_->IsWriting()) {
      channel_->EnableWriting();
    }
  }
}

bool TLSConnection::WillBlock(int ret, int *error_out) noexcept {
  VLOG(11) << "TLSConnection::WillBlock(ret=" << ret << ")";
  int error = *error_out = SSL_get_error(ssl_, ret);
  if (error == SSL_ERROR_WANT_READ) {
    // Register for read event if not already.
    channel_->EnableReading();
    return true;
  } else if (error == SSL_ERROR_WANT_WRITE) {
    // Register for write event if not already.
    channel_->EnableWriting();
    return true;
  } else {
    // SSL_ERROR_ZERO_RETURN is proccessed here so we can get some detail
    // in the log.
    long last_error = ERR_get_error();
    VLOG(11) << "TLSConnection(fd=" << socket_->fd() << ", state=" 
             << int(state_) << ", ssl_state=" << int(ssl_state_)
             << ", events=" << channel_->events() << "): SSL error: "
             << error << ", errno: " << ret << ", read: " 
             << BIO_number_read(SSL_get_rbio(ssl_)) << ", written: "
             << BIO_number_written(SSL_get_wbio(ssl_)) << ", func: "
             << ERR_func_error_string(last_error) << ", reason: " 
             << ERR_reason_error_string(last_error);
    if (error != SSL_ERROR_SYSCALL) {
      if (error == SSL_ERROR_SSL) {
        *error_out = last_error;
      }
      if ((unsigned long)last_error < 0x8000) {
        errno = ENOSYS;
      } else {
        errno = last_error;
      }
    }
    ERR_clear_error();
    return false;
  }
}

void TLSConnection::HandleAccept() {
  VLOG(11) << "TLSConnection::HandleAccept() this=" << this << ", fd=" 
           << socket_->fd() << ", state=" << int(state_) << ", ssl_state="
           << int(ssl_state_) << ", events=" << channel_->events();
  assert(state_ == StateE::kConnected && 
         ssl_state_ == SSLStateEnum::STATE_ACCEPTING);

  SSL_set_fd(ssl_, socket_->fd());

  errno = 0;
  int ret = SSL_accept(ssl_);
  if (ret <= 0) {
    int error;
    if (WillBlock(ret, &error)) {
      VLOG(11) << "TLSConnection::HandleAccept() SSL_accept() again";
      return;
    } else {
      ssl_state_ = STATE_ERROR;
      Shutdown();
      LOG(ERROR) << "SSL_accept error, close connection";
      return;
    }
  }

  ssl_state_ = SSLStateEnum::STATE_ESTABLISHED;
  connection_callback_(shared_from_this());
}

void TLSConnection::PerformRead() {
  int saved_errno = 0;
  ssize_t n = input_buffer_.ReadTls(ssl_, &saved_errno);
  if (n <= 0) {
    int error = SSL_get_error(ssl_, n);
    if (error == SSL_ERROR_WANT_READ) {
      // Not data read.
      VLOG(11) << "connection " << name_  << "not data recv";
    } else if (error == SSL_ERROR_WANT_WRITE) {
      LOG(ERROR) << "TLSConnection(fd=" << socket_->fd() << ", state=" 
                 << int(state_) << ", ssl_state=" << int(ssl_state_) 
                 << ", events=" << channel_->events() << "): " 
                 << "unsupported SSL renegotiation during read";
      ERR_clear_error();
    } else {
      long last_error = ERR_get_error();

      VLOG(11) << "TLSConnection(fd=" << socket_->fd() << ", state=" 
               << int(state_) << ", ssl_state=" << int(ssl_state_) 
               << ", events=" << channel_->events() << "): " 
               << "bytes: " << n << ", error: " << error << ", errno: "
               << errno << ", func: " << ERR_func_error_string(last_error) 
               << ", reason: " << ERR_reason_error_string(last_error);
      if (zero_return(error, n) || n == -1) {
        LOG(INFO) << "the connection " << name_ << " close";
        HandleClose();
      }
      ERR_clear_error();
    }
  } else {
    message_callback_(shared_from_this(), &input_buffer_);
  } 
}

void TLSConnection::PerformWrite() {
  if (channel_->IsWriting()) {
    ssize_t n = SSL_write(ssl_,
                          output_buffer_.Peek(),
                          output_buffer_.ReadableBytes());
    if (n <= 0) {
      long error = SSL_get_error(ssl_, n);
      if (error == SSL_ERROR_WANT_WRITE) {
        // Not data wrriten.
        ERR_clear_error();
        return ;
      } else if (error == SSL_ERROR_WANT_READ) {
        LOG(ERROR) << "TLSConnection(fd=" << socket_->fd() << ", state="
                   << int(state_) << ", ssl_state=" << int(ssl_state_) 
                   << ", events=" << channel_->events() << "): "
                   << "unsupported SSL renegotiation during write";
        ERR_clear_error();
      } else {
        long last_error = ERR_get_error();
        LOG(ERROR) << "TLSConnection(fd=" << socket_->fd() << ", state="
                   << int(state_) << ", ssl_state=" << int(ssl_state_)
                   << ", events=" << channel_->events() << "): SSL error: "
                   << error << ", errno: " << errno << ", func: "
                   << ERR_func_error_string(last_error) << ", reason: "
                   << ERR_reason_error_string(last_error);
        ERR_clear_error();
        if (!zero_return(error, n)) {
          HandleClose();
          return ;
        }
      }
    }

    if (n > 0) {
      output_buffer_.Retrieve(n);
      if (output_buffer_.ReadableBytes() == 0) {
        channel_->DisableWriting();
        if (state_ == kDisconnecting) {
          ShutdownInLoop();
        }
      } else {
        VLOG(11) << "I am going to write more data";
      }
    } else {
      LOG(ERROR) << "write error, res: " << n << " errno: " << errno 
                 << " ssl_err: " << SSL_get_error(ssl_, n);
    } // end if (n > 0)
  } else {
    VLOG(11) << "connection fd=" << channel_->fd() << " is down, no more writing";
  } // end if (channel_->IsWriting())
}

}  
}
